import os

import torch
from torch import nn
from yolo_v3_net import Yolo_V3_Net
from PIL import Image,ImageDraw
from utils import *
from dataset import *
from config import *

class_num={
    0:'person',
    1:'horse',
    2:'bicycle'
}
device = torch.device('cpu')
class Detector(nn.Module):
    def __init__(self):
        super(Detector, self).__init__()

        self.weights = 'params/netloss100.pt'
        self.net = Yolo_V3_Net().to(device)
        if os.path.exists(self.weights):
            self.net.load_state_dict(torch.load(self.weights))
            print('加载权重成功')
        # 调用 self.net.eval() 将模型 self.net 切换到评估模式。
        # 在评估模式下，一些特定的层，如 Dropout 和 Batch Normalization 层，会被固定，不再进行参数更新。
        # 这是因为在评估阶段，我们不需要计算梯度，而是只关注模型的输出。
        self.net.eval()

    def forward(self,input,thresh,anchors,case):
        output_13,output_26,output_52=self.net(input)
        index_13,bias_13=self.get_index_and_bias(output_13,thresh)
        boxes_13=self.get_true_position(index_13,bias_13,32,anchors[13],case)

        index_26, bias_26 = self.get_index_and_bias(output_26, thresh)
        boxes_26 = self.get_true_position(index_26, bias_26, 16, anchors[26],case)

        index_52, bias_52 = self.get_index_and_bias(output_52, thresh)
        boxes_52 = self.get_true_position(index_52, bias_52, 8, anchors[52],case)

        # 与 torch.stack() 不同的是，torch.cat() 不会增加维度，而是将输入张量的指定维度上的大小累加。
        return torch.cat([boxes_13,boxes_26,boxes_52],dim=0)

    # 获取object的索引和偏移量的函数
    # thresh，置信度筛选的阈值
    def get_index_and_bias(self, output, thresh):
        output = output.permute(0, 2, 3, 1)
        output = output.reshape(output.size(0), output.size(1), output.size(2), 3, -1)  # N,W,H,3,15

        mask = output[..., 0] > thresh  # N,H,W,3

        # index = mask.nonzero() 的目的是找到 mask 中为 True 的元素的索引。
        # nonzero() 是一个 PyTorch 中的方法，它会返回一个张量，其中包含输入张量中非零元素的索引。
        # 在这里，由于 mask 是一个布尔张量，实际上是找到 mask 中为 True 的元素所在的位置。
        # 返回的 index 张量是一个二维张量，每一行表示一个非零元素的索引。如果 mask 中有多个维度，那么 index 张量中每一列对应每个维度上的索引值。
        index = mask.nonzero() # [0, 12, 12, 2] [0, 12, 5, 2] [0, 12, 4, 1] [N W H 3]
        bias = output[mask]  # 后面的数量为15的张量的值

        return index, bias

    def get_true_position(self,index,bias,t,anchors,case):

        anchors=torch.Tensor(anchors)
        # 通过 a = index[:, 3] 的方式，将 index 张量的第四列（索引为3的列）提取出来，赋值给变量 a。
        # 冒号用于对索引变量 index 进行切片操作。它表示选择 index 张量的所有行，但只选择第四列（索引为3的列）的元素。
        # 冒号的一般语法是 start:stop:step，其中 start 是切片的开始位置，stop 是切片的结束位置（不包含在切片中），step 是步长。当某个值被省略时，它会取默认值，即从开头或末尾延伸。
        a=index[:,3]

        # t 为一个grid cell对应的边长的大小
        # case为该图片的缩放因子
        cx=(index[:,1].float()+bias[:,1].float())*t/case
        cy=(index[:,2].float()+bias[:,2].float())*t/case

        # 在存入标签值w,h值的时候，做了log的操作，现在这里需要还原做指数操作
        # a 表示的是哪个框 （0，1，2）三个值中的索引
        w=anchors[a,0]*torch.exp(bias[:,3])/case
        h=anchors[a,1]*torch.exp(bias[:,4])/case
        # 置信度为最后一个维度上的第0个元素
        p=bias[:,0]
        # One-hot编码
        cls_p=bias[:,5:]
        # One-hot编码的下标
        cls_index=torch.argmax(cls_p,dim=1)
        # 之前在训练的时候，对置信度p的值做了sigmod处理，这里还是需要对置信度的值进行sigmod处理
        # 将 torch.sigmoid(p)、cx、cy、w、h、cls_index 这些张量在维度 1 上进行拼接。
        # 注意，维度的编号从 0 开始。因此，如果 dim=1，则表示在第二个维度上进行拼接。拼接后的张量将具有拼接维度上增加的维度。
        return torch.stack([torch.sigmoid(p),cx,cy,w,h,cls_index],dim=1)

def inverse_sigmoid(y):
    return math.log(y / (1 - y))

if __name__ == '__main__':
    detector=Detector()
    img=Image.open('images/15.jpg')
    _img=make_416_image('images/15.jpg')
    temp=max(_img.size)
    case=416/temp
    _img=_img.resize((416,416))
    _img=tf(_img).to(device)
    # torch.unsqueeze() 函数用于在指定维度上给张量增加一个维度。
    # torch.unsqueeze() 函数接受一个张量作为输入，并在指定的维度上增加一个维度。这个函数的作用类似于在相应的位置上插入一个大小为 1 的维度。
    # 在这个例子中，_img 是一个张量，而 dim=0 表示在第一个维度上增加一个维度。
    # 假设 _img 的形状为 (C, H, W)，其中 C 表示通道数，H 表示高度，W 表示宽度。
    # 那么，torch.unsqueeze(_img, dim=0) 的结果将具有形状 (1, C, H, W)。在该结果中，第一个维度上增加了一个维度，大小为 1，而后续维度保持不变。
    # 这种操作常用于将单个样本或图像扩展为批次样本或批次图像的形式，以便于输入深度学习模型进行批处理。
    _img=torch.unsqueeze(_img,dim=0)
    results=detector(_img,inverse_sigmoid(0.1),ANCHORS,case)

    draw=ImageDraw.Draw(img)

    for rst in results:
        x1=rst[1]-0.5*rst[3]
        x2=rst[1]+0.5*rst[3]
        y1=rst[2]-0.5*rst[4]
        y2=rst[2]+0.5*rst[4]
        print(x1,y1,x2,y2)
        print('class',class_num[int(rst[5])])
        # 在框的左上角坐标处，写类别的信息，和置信度的信息
        # rst[0].item() 函数将列表 rst 中的第一个元素转换为 Python 的标量值。
        # str(rst[0].item())[:4] 将该标量值转换为字符串，并截取字符串的前四个字符。
        draw.text((x1,y1),str(class_num[int(rst[5].item())])+str(rst[0].item())[:4])
        print(str(rst[0].item())[:4]);
        draw.rectangle((x1,y1,x2,y2),outline='red',width=1)
    img.show()
